import numpy as np
from tensorboardX import SummaryWriter
import os
import nibabel as nib
from PIL import Image

#  定义.nii转换为.jpg的方法
def Nii2jpg_x255(nifti_path, output_dir, file):
    # 加载 NIfTI 文件
    img = nib.load(nifti_path)
    n1 = img.get_fdata()

    # 检查 n1 是否为二维数组
    if n1.ndim != 2:
        raise ValueError("The loaded NIfTI data is not a 2D array as expected.")

    # 归一化数据到 0-255 范围并转换为 uint8
    normalized_data = ((n1 - n1.min()) / (n1.max() - n1.min()) * 255).astype(np.uint8)

    # 创建输出文件路径
    file_jpg = file.replace(".nii", ".jpg")
    output_path = os.path.join(output_dir, file_jpg)

    # 保存为 JPG 图像
    Image.fromarray(normalized_data).save(output_path)


# 待转换的各.nii文件路径
data_trbase_path = 'D:/DeepModel/le-nii/exing'
data_tebase_path = 'D:/DeepModel/le-nii/lxing'

# 各.jpg文件路径
data_trpath = 'D:\DeepModel\le_jpg\lxing'
data_tepath = 'D:\DeepModel\le_jpg\exing'


# 遍历各个根文件夹中的每个NIfTI文件，并将其转换为JPG图像
for root, dirs, files in os.walk(data_trbase_path):
    for file in files:
        if file.endswith('.nii'):
            nifti_path = os.path.join(root, file)
            # save_nifti_as_jpg(nifti_path, train_lpath)
            Nii2jpg_x255(nifti_path, data_tepath, file)
for root, dirs, files in os.walk(data_tebase_path):
    for file in files:
        if file.endswith('.nii'):
            nifti_path = os.path.join(root, file)
            # save_nifti_as_jpg(nifti_path, train_epath)
            Nii2jpg_x255(nifti_path, data_trpath, file)
